#pragma once
#ifndef ZQ_TOOLS_H_
#define ZQ_TOOLS_H_

#include <algorithm>
#include <vector>
#include <cnpy.h>
#include "stream_tools.h"

template<unsigned N_CH, unsigned BIT>
void load_image(const cnpy::NpyArray& img, data_stream<N_CH * BIT>& out)
{   
    assert(img.word_size == sizeof(int32_t));
    assert(img.shape.size() == 3);
    assert(img.shape[2] == N_CH);
    // assert(out.empty());

    const int32_t* ptr = img.data<int32_t>();
    const size_t iters = img.shape[0] * img.shape[1];
    for (size_t i = 0; i < iters; ++i)
    {
        ap_uint<N_CH * BIT> buf;
        for (unsigned j = 0; j < N_CH; ++j)
        {
            ap_uint<BIT> x = ptr[j];
            buf(SLICE(BIT, j)) = x;
        }
        out.write(buf);
        ptr += N_CH;
    }
    return;
};

template<
    unsigned N_IO,     // # of IO
    unsigned N_CH,     // # of CHannel
    unsigned BIT,      // nBit
    typename T = int32_t
>
void load_DW_weight(const cnpy::NpyArray& npy, ap_uint<N_IO * BIT> dst[N_CH / N_IO][9])
{
    static_assert(N_CH >= N_IO, "load_DW_weight");
    static_assert(N_CH % N_IO == 0, "load_DW_weight");
    constexpr unsigned FOLD = N_CH / N_IO;

    assert(npy.word_size == sizeof(T));
    assert(npy.shape.size() == 3);
    assert(npy.shape[0] == 3);
    assert(npy.shape[1] == 3);
    assert(npy.shape[2] == N_CH);

    for (size_t f = 0; f < FOLD; ++f)
    {
        for (size_t k = 0; k < 9; ++k)
        {
            const T* src = npy.data<T>() + N_CH * k + FOLD * f;
            ap_uint<N_IO * BIT> buf;
            for (size_t i = 0; i < N_IO; ++i)
            {
                ap_int<BIT> temp = src[i];
                buf(SLICE(BIT, i)) = temp;
            }
            dst[f][k] = buf;
        }
    }

    return;
}

template<
    unsigned N_IN,     // # of INput
    unsigned N_OUT,     // # of OUTput
    unsigned N_ICH,     // # of Input CHannel
    unsigned N_OCH,     // # of Output CHannel
    unsigned BIT,      // nBit
    typename T = int32_t
>
void load_PW_weight(const cnpy::NpyArray& npy, ap_uint<N_OUT * N_IN * BIT> dst[N_OCH / N_OUT][N_ICH / N_IN])
{
    static_assert(N_ICH >= N_IN,  "load_PW_weight");
    static_assert(N_OCH >= N_OUT, "load_PW_weight");
    static_assert(N_ICH % N_IN  == 0, "load_PW_weight");
    static_assert(N_OCH % N_OUT == 0, "load_PW_weight");
    constexpr unsigned FOLD_I = N_ICH / N_IN;
    constexpr unsigned FOLD_O = N_OCH / N_OUT;

    assert(npy.word_size == sizeof(T));
    assert(npy.shape.size() == 2);
    assert(npy.shape[0] == N_OCH);
    assert(npy.shape[1] == N_ICH);

    const T* src = npy.data<T>();
    for (size_t fo = 0; fo < FOLD_O; ++fo)
    {
        for (size_t fi = 0; fi < FOLD_I; ++fi)
        {
            ap_uint<N_OUT * N_IN * BIT> buf;
            for (size_t o = 0; o < N_OUT; ++o)
            {
                for (size_t i = 0; i < N_IN; ++i)
                {
                    size_t idx = N_OUT * N_ICH * fo + N_ICH * o + N_IN * fi + i;
                    ap_int<BIT> temp = src[idx];
                    buf(SLICE(BIT, N_IN * o + i)) = temp;
                }
            }
            dst[fo][fi] = buf;
        }
    }

    return;
}

template<
    unsigned N_IO,     // # of IO
    unsigned BIT,      // nBit
    typename T = int32_t
>
void load_inOrder(const T* src, ap_uint<N_IO * BIT>* dst, size_t len_dst)
{
    for (size_t i = 0; i < len_dst; ++i)
    {
        ap_uint<N_IO * BIT> buf;
        for (size_t j = 0; j < N_IO; ++j)
        {
            ap_int<BIT> temp = *(src++);
            buf(SLICE(BIT, j)) = temp;
        }
        *(dst++) = buf;
    }
    return;
}

template<
    unsigned N_IO,     // # of IO
    unsigned N_CH,     // # of CHannel
    unsigned BIT,      // nBit
    typename T = int32_t
>
void load_array1D(const cnpy::NpyArray& npy, ap_uint<N_IO * BIT> dst[N_CH / N_IO])
{
    static_assert(N_CH >= N_IO, "load_array1D");
    static_assert(N_CH % N_IO == 0, "load_array1D");
    constexpr unsigned FOLD = N_CH / N_IO;

    assert(npy.word_size == sizeof(T));
    assert(npy.shape.size() == 1);
    assert(npy.shape[0] == N_CH);

    load_inOrder<N_IO, BIT, T>(npy.data<T>(), dst, FOLD);
    return;
}

template<unsigned N_CH, unsigned BIT, typename T = int32_t, typename U = ap_uint<BIT>>
bool check(data_stream<N_CH * BIT>& io, const cnpy::NpyArray& gold, unsigned nbatch = 1)
{
    assert(gold.word_size == sizeof(T));
    assert(gold.shape.size() == 3);
    assert(gold.shape[0] * gold.shape[1] * gold.shape[2] % N_CH == 0);
    const size_t vec_num = gold.shape[0] * gold.shape[1] * gold.shape[2] / N_CH;
    assert(nbatch * vec_num == io.size());

    std::vector<T> buf;
    buf.reserve(N_CH);

    for (unsigned b = 0; b < nbatch; ++b)
    {
        const T* ptr = gold.data<T>();
        for (size_t i = 0; i < vec_num; ++i)
        {
            const ap_uint<N_CH * BIT> temp = io.read();
            for (size_t j = 0; j < N_CH; ++j)
            {
                U x = temp(SLICE(BIT, j));
                buf.push_back(x);
            }
            io.write(temp);
            bool pass = std::equal(buf.begin(), buf.end(), ptr);
            if (!pass)
            {
                printf("Boom @ #%u %ld\n", b, i);
                return false;
            }
            ptr += N_CH;
            buf.clear();
        }
    }
    return true;
};

template<unsigned N_CH, unsigned BIT>
void save_stream(const std::string& fname, data_stream<N_CH * BIT>& io)
{
    const size_t vec_len = io.size();
    std::vector<size_t> shape = {vec_len, N_CH};
    std::vector<int> data;
    data.reserve(vec_len * N_CH);
    for (size_t i = 0; i < vec_len; ++i)
    {
        ap_uint<N_CH * BIT> temp = io.read();;
        for (size_t j = 0; j < N_CH; ++j)
        {
            int x = temp(SLICE(BIT, j));
            data.push_back(x);
        }
        io.write(temp);
    }
    cnpy::npy_save(fname, data.data(), shape);
    return;
};

template<unsigned WIDTH>
void clear_stream(data_stream<WIDTH>& in)
{
    while(!in.empty())
        in.read();
    return;
};

#endif
